#pragma once
#include <cassert>

namespace UF6
{
	class UnionFind
	{
	public:
		UnionFind(int n)
		{
			parent = new int[n];
			rank = new int[n];
			count = n;

			for (size_t i = 0; i < count; i++)
			{
				parent[i] = i;
				rank[i] = 1;
			}
		}

		~UnionFind()
		{
			delete[] parent;
			delete[] rank;
		}

		int Find(int idx)
		{
			assert(idx >= 0 && idx < count);
			if (idx != parent[idx])
			{
				parent[idx] = Find(parent[idx]);
			}
			return parent[idx];
		}

		bool IsConnected(int p, int q)
		{
			assert(p >= 0 && p < count);
			assert(q >= 0 && q < count);

			return Find(p) == Find(q);
		}

		void Union(int p, int q)
		{
			int pID = Find(p);
			int qID = Find(q);

			if (pID == qID)
				return;

			if (rank[pID] < rank[qID])
			{
				parent[pID] = qID;
			}
			else if (rank[pID] > rank[qID])
			{
				parent[qID] = pID;
			}
			else
			{
				parent[pID] = qID;
				rank[qID]++;
			}
		}

	private:
		int* parent;
		int* rank;
		int count;
	};
}
